import numpy as np
import os
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

workspace = 'C:\\Users\\Administrator\\Desktop'
filename = '轨迹数据.txt'
fig = plt.figure()
#ax = fig.add_subplot(111, projection='3d')


"=============================[轨迹插补]==========================="


def read(file_dir):
    postion_x = []
    postion_y = []
    postion_z = []
    rot_n_x = []
    rot_n_y = []
    rot_n_z = []
    rot_o_x = []
    rot_o_y = []
    rot_o_z = []
    rot_v_x = []
    rot_v_y = []
    rot_v_z = []
    for line in open(file_dir, encoding='utf-8'):
        str = "".join([c for c in line if c in '0123456789.-,'])
        #print(str)
        nums = [float(num) for num in str.split(",")]
        # 姿态n坐标
        rot_n_x.append(nums[0])
        rot_n_y.append(nums[1])
        rot_n_z.append(nums[2])
        # 姿态o坐标
        rot_o_x.append(nums[3])
        rot_o_y.append(nums[4])
        rot_o_z.append(nums[5])
        # 姿态v坐标
        rot_v_x.append(nums[6])
        rot_v_y.append(nums[7])
        rot_v_z.append(nums[8])
        # 位置坐标
        postion_x.append(nums[9])
        postion_y.append(nums[10])
        postion_z.append(nums[11])
    return (rot_n_x, rot_n_y, rot_n_z, rot_o_x, rot_o_y, rot_o_z, rot_v_x, rot_v_y, rot_v_z, postion_x, postion_y, postion_z)

def drawPath(path):
    ax.scatter3D(path[9], path[10], path[11], c='k', marker='.')

def drawPose(path):
    path = np.array(path).transpose()
    for point in path:
        #print(point)
        x =  point[9]
        y =  point[10]
        z =  point[11]
        k = 10
        x1 = k * point[0]
        y1 = k * point[1]
        z1 = k * point[2]
        ax.scatter3D(x+x1, y+y1, z+z1, c='g', marker='.')
        plt.plot([x+x1,x], [y+y1,y], [z+z1,z], '-g')

        x2 = k * point[3]
        y2 = k * point[4]
        z2 = k * point[5]
        ax.scatter3D(x+x2, y+y2, z+z2, c='r', marker='.')
        plt.plot([x+x2,x], [y+y2,y], [z+z2,z], '-r')

        x3 = k * point[6]
        y3 = k * point[7]
        z3 = k * point[8]
        ax.scatter3D(x+x3, y+y3, z+z3, c='b', marker='.')
        plt.plot([x+x3,x], [y+y3,y], [z+z3,z], '-b')

def draw(path):
    mid = len(path[0]) // 2
    ax.set_xlim3d(path[9][mid]  - 50, path[9][mid] + 50)
    ax.set_ylim3d(path[10][mid] - 50, path[10][mid] + 50)
    ax.set_zlim3d(path[11][mid] - 50, path[11][mid] + 50)
    plt.show()


def draw_pose_path_from(file_dir):
    # 读取文件
    path = read(file_dir)
    # 绘轨迹
    drawPath(path)
    # 绘姿态
    drawPose(path)
    # 定位到视图中间
    draw(path)


"=============================[定点转]==========================="


def draw_a_pose(matrix):
    k = 100
    n = [k*matrix[0][0], k*matrix[1][0], k*matrix[2][0]]
    o = [k*matrix[0][1], k*matrix[1][1], k*matrix[2][1]]
    a = [k*matrix[0][2], k*matrix[1][2], k*matrix[2][2]]
    x = matrix[0][3]
    y = matrix[1][3]
    z = matrix[2][3]

    ax.scatter3D(x, y, z, c='k', marker='.')
    ax.scatter3D(x + n[0], y + n[1], z + n[2], c='g', marker='.')
    ax.scatter3D(x + o[0], y + o[1], z + o[2], c='r', marker='.')
    ax.scatter3D(x + a[0], y + a[1], z + a[2], c='b', marker='.')
    plt.plot([x + n[0], x], [y + n[1], y], [z + n[2], z], '-g')
    plt.plot([x + o[0], x], [y + o[1], y], [z + o[2], z], '-r')
    plt.plot([x + a[0], x], [y + a[1], y], [z + a[2], z], '-b')



def draw_in_middle(mat):
    x = mat[0][3]
    y = mat[1][3]
    z = mat[2][3]
    ax.set_xlim3d(x - 50, x + 50)
    ax.set_ylim3d(y - 50, x + 50)
    ax.set_zlim3d(y - 50, x + 50)
    plt.show()

def fix_point_rotate(cur):
    curMatrix = np.array([cur[0], cur[3], cur[6], cur[9],
                          cur[1], cur[4], cur[7], cur[10],
                          cur[2], cur[5], cur[8], cur[11],
                          0, 0, 0, 1]).reshape(4, 4)
    # 10°插补一次
    theta = 20/180*np.pi
    cnt = int(2 * np.pi / theta)
    rotx = np.array([1, 0, 0, 0,
                     0, np.cos(theta), -np.sin(theta), 0,
                     0, np.sin(theta), np.cos(theta), 0,
                     0, 0, 0, 1]).reshape(4,4)

    roty = np.array([np.cos(theta), 0, np.sin(theta), 0,
                     0, 1, 0, 0,
                     -np.sin(theta), 0, np.cos(theta), 0,
                     0, 0, 0, 1]).reshape(4, 4)

    rotz = np.array([np.cos(theta), -np.sin(theta), 0, 0,
                     np.sin(theta), np.cos(theta), 0, 0,
                     0, 0, 1, 0,
                     0, 0, 0, 1]).reshape(4,4)
    print(curMatrix)
    print(rotx)
    print("cnt:", cnt)
    i = 0
    while i < cnt:
        #print("cur:", curMatrix)
        targetMatrix = np.dot(curMatrix, roty)
        #print("tar:", targetMatrix)
        draw_a_pose(targetMatrix)
        i += 1
        curMatrix = targetMatrix

    draw_in_middle(curMatrix)


def readSpeedData(file_dir):
    times = []
    speeds = []
    trans = []
    for line in open(file_dir, encoding='utf-8'):
        if not line[0].isdigit():
            continue
        str = "".join([c for c in line if c in '0123456789.-,*'])
        arr = [float(val) for val in str.split(',')]
        #print(arr)
        times.append(arr[0])
        speeds.append(arr[1])
        if arr[2] == 1:
            trans.append((arr[0], arr[1]))
    #print(trans)
    return times, speeds, trans


def drawSpeedCerve(times, speeds, trans):
    # 画过渡点
    for tran in trans:
        plt.plot(tran[0], tran[1], 'ro')
    plt.plot(times, speeds)
    plt.title('x - v')
    plt.ylabel('speed')
    plt.xlabel('distance')
    plt.show()
    #print(speeds)

"============================[空间三点计算圆心]===================================="
def getCenterOfCircle(points):
	x1 = points[0][0]
	x2 = points[1][0]
	x3 = points[2][0]
	y1 = points[0][1]
	y2 = points[1][1]
	y3 = points[2][1]
	z1 = points[0][2]
	z2 = points[1][2]
	z3 = points[2][2]

	a1 = (y1*z2 - y2*z1 - y1*z3 + y3*z1 + y2*z3 - y3*z2)
	b1 = -(x1*z2 - x2*z1 - x1*z3 + x3*z1 + x2*z3 - x3*z2)
	c1 = (x1*y2 - x2*y1 - x1*y3 + x3*y1 + x2*y3 - x3*y2)
	d1 = -(x1*y2*z3 - x1*y3*z2 - x2*y1*z3 + x2*y3*z1 + x3*y1*z2 - x3*y2*z1)

	a2 = 2 * (x2 - x1)
	b2 = 2 * (y2 - y1)
	c2 = 2 * (z2 - z1)
	d2 = x1*x1 + y1*y1 + z1*z1 - x2*x2 - y2*y2 - z2*z2

	a3 = 2 * (x3 - x1)
	b3 = 2 * (y3 - y1)
	c3 = 2 * (z3 - z1)
	d3 = x1*x1 + y1*y1 + z1*z1 - x3*x3 - y3*y3 - z3*z3

	cx = -(b1*c2*d3 - b1*c3*d2 - b2*c1*d3 + b2*c3*d1 + b3*c1*d2 - b3*c2*d1)/(a1*b2*c3 - a1*b3*c2 - a2*b1*c3 + a2*b3*c1 + a3*b1*c2 - a3*b2*c1)
	cy =  (a1*c2*d3 - a1*c3*d2 - a2*c1*d3 + a2*c3*d1 + a3*c1*d2 - a3*c2*d1)/(a1*b2*c3 - a1*b3*c2 - a2*b1*c3 + a2*b3*c1 + a3*b1*c2 - a3*b2*c1)
	cz = -(a1*b2*d3 - a1*b3*d2 - a2*b1*d3 + a2*b3*d1 + a3*b1*d2 - a3*b2*d1)/(a1*b2*c3 - a1*b3*c2 - a2*b1*c3 + a2*b3*c1 + a3*b1*c2 - a3*b2*c1)

	return [cx, cy, cz]

def draw_a_point(point):
    x = point[0]
    y = point[1]
    z = point[2]
    ax.scatter3D(x, y, z, c='k', marker='.')

def draw_in_middle_with(point):
    x = point[0]
    y = point[1]
    z = point[2]
    ax.set_xlim3d(x - 5, x + 5)
    ax.set_ylim3d(y - 5, x + 5)
    ax.set_zlim3d(y - 5, x + 5)
    plt.show()


"============================[主函数]===================================="

if __name__ == '__main__':
    # 绘制机器人轨迹及姿态
    #draw_pose_path_from(os.path.join(workspace, filename))

    # 定点转
    #curPose = [-0.18, -0.03, 0.98, 0.14, -0.99, 0.00, 0.97, 0.14, 0.18, 437.63, 58.41, 400.38]
    #fix_point_rotate(curPose)

    # 速度仿真
    time, speeds, trans = readSpeedData(os.path.join(workspace, filename))
    drawSpeedCerve(time, speeds, trans)


    # 空间三点计算圆心
    # points = [[1,5,1], [2,-1,4], [1,-2,1]]
    # draw_a_point(points[0])
    # draw_a_point(points[1])
    # draw_a_point(points[2])
    # center = getCenterOfCircle(points)
    # draw_a_point(center)
    # draw_in_middle_with(center)
